import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from torchvision.utils import save_image
import os

# 定义生成器网络结构（必须与训练时相同）
class Generator(nn.Module):
    def __init__(self, latent_dim, img_shape):
        super(Generator, self).__init__()
        self.img_shape = img_shape

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Linear(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *block(latent_dim, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Linear(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def forward(self, z):
        img = self.model(z)
        img = img.view(img.size(0), *self.img_shape)
        return img

# 参数设置（必须与训练时相同）
latent_dim = 100
img_shape = (1, 28, 28)
model_path = './gan_models'

# 初始化生成器
generator = Generator(latent_dim, img_shape)

# 加载训练好的模型
# 修改加载方式，从保存的字典中提取模型状态
try:
    generator_checkpoint = torch.load(os.path.join(model_path, 'generator.pth'))
    generator.load_state_dict(generator_checkpoint['model_state_dict'])
    print("成功加载生成器模型")
except Exception as e:
    print(f"加载模型失败: {str(e)}")
    exit(1)

generator.eval()  # 设置为评估模式

# 生成测试图像
with torch.no_grad():
    # 生成随机噪声
    z = torch.randn(1, latent_dim)

    # 生成图像
    gen_img = generator(z)

    # 转换为numpy数组并显示
    img = gen_img[0].permute(1, 2, 0).numpy()
    plt.imshow(img[:, :, 0], cmap='gray')
    plt.axis('off')
    plt.show()

    # 保存生成的图像
    save_image(gen_img.data, "./gan_models/generated_image.png", nrow=1, normalize=True)
